package com.nd.arcSoft.service.impl;

import com.arcsoft.face.FaceFeature;
import com.nd.arcSoft.entity.FaceUserInfo;
import com.nd.arcSoft.entity.RespGrabber;
import com.nd.arcSoft.entity.RespStopCondition;
import com.nd.arcSoft.service.IArcSoftService;
import com.nd.arcSoft.service.IRespGrabberService;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.ffmpeg.global.avutil;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.FrameGrabber;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.springframework.scheduling.annotation.Async;
import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * @Author 张寒轩
 * @Date 2021/3/20 15:21
 * @Description 流媒体抓取实现
 */
@Slf4j
@Service
public class RespGrabberServiceImpl implements IRespGrabberService {

    private ExecutorService executorService;

    @Resource
    private IArcSoftService arcSoftService;

    @PostConstruct
    private void init() {
        int threadPoolSize = Runtime.getRuntime().availableProcessors();
        ThreadFactory springThreadFactory = new CustomizableThreadFactory("arcSoft-resp-");
        /**
         * 核心线程数=CPU核心数
         * 最大线程数=CPU核心数
         * 回收线程时间=0
         * 回收线程单位=秒
         * 队列=直接提交
         * 线程工厂=自定义工厂(设置线程名)
         * 拒绝策略=自定义
         */
        executorService = new ThreadPoolExecutor(
                threadPoolSize,
                threadPoolSize,
                0L,
                TimeUnit.SECONDS,
                new SynchronousQueue(),
                springThreadFactory,
                new MyRejectedExecutionHandler());
    }

    /**
     * @auther 张寒轩
     * @date 2021/3/31 16:51
     * @description 自定义拒绝方案，超出线程直接拒绝
     */
    private class MyRejectedExecutionHandler implements RejectedExecutionHandler{

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            throw new RuntimeException("超出RESP允许最大线程，最大线程为"+executor.getMaximumPoolSize());
        }
    }

    /**
     * @auther 张寒轩
     * @date 2021/3/25 14:08
     * @description resp帧抓取比对执行任务
     */
    private class RespGrabberTask implements Runnable {
        private RespGrabber respGrabber;
        private List<FaceUserInfo> faceUserInfos;
        private String filePath;
        private RespStopCondition stopCondition;
        private Consumer<List<FaceUserInfo>> consumer;

        public RespGrabberTask(RespGrabber respGrabber,
                               List<FaceUserInfo> faceUserInfos,
                               String filePath,
                               RespStopCondition stopCondition,
                               Consumer<List<FaceUserInfo>> consumer) {
            this.respGrabber = respGrabber;
            this.faceUserInfos = faceUserInfos;
            this.filePath = filePath;
            this.stopCondition = stopCondition;
            this.consumer = consumer;
        }

        @Override
        public void run() {
            //帧抓取工具
            FFmpegFrameGrabber grabber = createGrabber(respGrabber);
            Java2DFrameConverter java2DFrameConverter = new Java2DFrameConverter();
            Double frameRate = grabber.getFrameRate();
            int frameRateInt = frameRate.intValue();
            //计数器
            long num = 0;
            //帧
            Frame frame = null;
            String name = Thread.currentThread().getName();
            //循环执行
            while (true) {
                try {
                    log.debug("当前线程：{}，resp帧比对，第：{} 次开始", name, ++num);
                    Long startTime = System.currentTimeMillis();
                    //帧坐标
                    int index = 0;
                    loop1:
                    while (true) {
                        ++index;
                        frame = grabber.grabImage();
                        if (null != frame && index % frameRateInt == 0) {
                            break loop1;
                        }
                    }
                    BufferedImage bi = java2DFrameConverter.getBufferedImage(frame);
                    byte[] jpegs = imageToBytes(bi, "jpeg");
                    //保存截取到的图片
                    if (!StringUtils.isEmpty(filePath)) {
                        FileOutputStream outputStream = new FileOutputStream(filePath + name + "-" + num + ".jpeg");
                        ImageIO.write(bi, "jpeg", outputStream);
                        outputStream.close();
                    }
                    List<FaceFeature> faceFeatures = arcSoftService.extractFaceFeature(jpegs);
                    log.debug("当前线程：{}，resp帧比对，第：{} 次，匹配人脸数：{}", name, num, faceFeatures.size());
                    if (!CollectionUtils.isEmpty(faceFeatures)) {
                        for (int i = 0; i < faceFeatures.size(); i++) {
                            log.debug("当前线程：{}，resp帧比对，第：{} 次，开始第{}张人脸查找", name, num, i + 1);
                            List<FaceUserInfo> compareFaces = arcSoftService.compareFaceFeature(faceFeatures.get(i), faceUserInfos, 500);
                            log.debug("当前线程：{}，resp帧比对，第：{} 次，结束第{}张人脸查找,匹配人脸数：{}", name, num, i + 1, compareFaces.size());
                            //匹配到的人脸执行业务处理
                            consumer.accept(compareFaces);
                        }
                    }
                    log.debug("当前线程：{}，resp帧比对，第：{} 次结束,耗时：{}毫秒", name, num, System.currentTimeMillis() - startTime);
                    //判断是否结束
                    if (null != stopCondition) {
                        if (null != stopCondition.getEndNum() && num >= stopCondition.getEndNum()) {
                            grabber.stop();
                            frame.close();
                            return;
                        }
                        if (null != stopCondition.getEndTime() && startTime >= stopCondition.getEndTime()) {
                            grabber.stop();
                            frame.close();
                            return;
                        }
                    }
                } catch (Exception e) {
                    log.error("", e);
                    if (grabber != null) {
                        try {
                            grabber.stop();
                            frame.close();
                        } catch (FrameGrabber.Exception ex) {
                            log.error("", e);
                        } finally {
                            return;
                        }
                    }
                }
            }
        }
    }

    /**
     * @param respGrabber resp配置
     * @return FFmpegFrameGrabber resp连接
     * @auther 张寒轩
     * @date 2021/3/20 15:41
     * @description 创建resp连接
     */
    private FFmpegFrameGrabber createGrabber(RespGrabber respGrabber) {
        try {
            FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(respGrabber.getUrl());
            grabber.setFrameRate(respGrabber.getFrameRate());
            grabber.setVideoBitrate(respGrabber.getBitRate());
            grabber.setImageWidth(respGrabber.getFrameWidth());
            grabber.setImageHeight(respGrabber.getFrameHeight());
            //grabber.setOption("rtsp_transport", "tcp");
            //过滤警告
            avutil.av_log_set_level(avutil.AV_LOG_ERROR);
            grabber.start();
            return grabber;
        } catch (Exception e) {
            log.error("", e);
        }
        return null;
    }

    /**
     * @param bi     图片缓存
     * @param format 图片格式
     * @return byte
     * @auther 张寒轩
     * @date 2021/3/20 16:15
     * @description 返回图片数组
     */
    private byte[] imageToBytes(BufferedImage bi, String format) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try {
            ImageIO.write(bi, format, baos);
        } catch (IOException e) {
            log.error(e.getMessage());
            return null;
        }
        return baos.toByteArray();
    }


    /**
     * @param respGrabber   resp连接对象
     * @param faceUserInfos 人脸池
     * @param filePath      若需要保存每一帧图片，传递的路径
     * @param stopCondition 结束条件
     * @param consumer      匹配到的人脸处理函数
     * @return void
     * @auther 张寒轩
     * @date 2021/3/20 17:54
     * @description 开启resp 图像处理
     */
    @Async
    @Override
    public void startGrabber(RespGrabber respGrabber, List<FaceUserInfo> faceUserInfos, String filePath, RespStopCondition stopCondition, Consumer<List<FaceUserInfo>> consumer) {
        //帧抓取工具
        FFmpegFrameGrabber grabber = createGrabber(respGrabber);
        Java2DFrameConverter java2DFrameConverter = new Java2DFrameConverter();
        Double frameRate = grabber.getFrameRate();
        int frameRateInt = frameRate.intValue();
        //计数器
        long num = 0;
        //帧
        Frame frame = null;
        String threadName = Thread.currentThread().getName();
        //循环执行
        while (true) {
            try {
                log.debug("线程：{}，resp帧比对，第：{} 次开始", threadName, ++num);
                Long startTime = System.currentTimeMillis();
                //帧坐标
                int index = 0;
                loop1:
                while (true) {
                    ++index;
                    frame = grabber.grabImage();
                    if (null != frame && index % frameRateInt == 0) {
                        break loop1;
                    }
                }
                BufferedImage bi = java2DFrameConverter.getBufferedImage(frame);
                byte[] jpegs = imageToBytes(bi, "jpeg");
                //保存截取到的图片
                if (!StringUtils.isEmpty(filePath)) {
                    FileOutputStream outputStream = new FileOutputStream(filePath + threadName + "-" + num + ".jpeg");
                    ImageIO.write(bi, "jpeg", outputStream);
                    outputStream.close();
                }
                List<FaceFeature> faceFeatures = arcSoftService.extractFaceFeature(jpegs);
                log.debug("线程：{}，resp帧比对，第：{} 次，匹配人脸数：{}", threadName, num, faceFeatures.size());
                if (!CollectionUtils.isEmpty(faceFeatures)) {
                    for (int i = 0; i < faceFeatures.size(); i++) {
                        log.debug("线程：{}，resp帧比对，第：{} 次，开始第{}张人脸查找", threadName, num, i + 1);
                        List<FaceUserInfo> compareFaces = arcSoftService.compareFaceFeature(faceFeatures.get(i), faceUserInfos, 100);
                        log.debug("线程：{}，resp帧比对，第：{} 次，结束第{}张人脸查找,匹配人脸数：{}", threadName, num, i + 1, compareFaces.size());
                        //匹配到的人脸执行业务处理
                        consumer.accept(compareFaces);
                    }
                }
                log.debug("线程：{}，resp帧比对，第：{} 次结束,耗时：{}毫秒", threadName, num, System.currentTimeMillis() - startTime);
                //判断是否结束
                if (null != stopCondition) {
                    if (null != stopCondition.getEndNum() && num >= stopCondition.getEndNum()) {
                        grabber.stop();
                        frame.close();
                        return;
                    }
                    if (null != stopCondition.getEndTime() && startTime >= stopCondition.getEndTime()) {
                        grabber.stop();
                        frame.close();
                        return;
                    }
                }
            } catch (Exception e) {
                log.error("", e);
                if (grabber != null) {
                    try {
                        grabber.stop();
                        frame.close();
                    } catch (FrameGrabber.Exception ex) {
                        log.error("", e);
                    }
                }
            }
        }

    }

    /**
     * @param respGrabbers  resp连接对象集合
     * @param faceUserInfos 人脸池
     * @param filePath      若需要保存每一帧图片，传递的路径
     * @param stopCondition 停止条件
     * @param consumer      匹配到的人脸处理函数
     * @return void
     * @auther 张寒轩
     * @date 2021/3/20 17:54
     * @description 开启resp 图像处理
     */
    @Async
    @Override
    public void startGrabber(List<RespGrabber> respGrabbers,
                             List<FaceUserInfo> faceUserInfos,
                             String filePath,
                             RespStopCondition stopCondition,
                             Consumer<List<FaceUserInfo>> consumer){

        for (int i = 0; i < respGrabbers.size(); i++) {
            executorService.execute(new RespGrabberServiceImpl.RespGrabberTask
                    (respGrabbers.get(i), faceUserInfos, filePath, stopCondition, consumer));
        }
    }

}